#include "kernel_operator.h"
using namespace AscendC;

class KernelAss
{
public:
    __aicore__ inline KernelAss() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR size, GM_ADDR stride, GM_ADDR storage_offset,
                                GM_ADDR y, uint32_t size_n, uint32_t blockLength)
    {
        xGm.SetGlobalBuffer((__gm__ DTYPE_X *)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ DTYPE_X *)y, blockLength);
        sizeGm.SetGlobalBuffer((__gm__ DTYPE_SIZE *)size, size_n);
        strideGm.SetGlobalBuffer((__gm__ DTYPE_SIZE *)stride, size_n);
        int32_t step1 = strideGm.GetValue(size_n - 1);
        int32_t offset = storage_offset[0];
        int32_t rpt1 = sizeGm.GetValue(size_n - 1);
        if (size_n == 1)
        {
            for (int i = 0; i < rpt1; i++)
            {
                yGm.SetValue(i, xGm.GetValue(offset + i * step1));
            }
        }
        else if (size_n == 2)
        {
            int32_t step2 = strideGm.GetValue(size_n - 2);
            int32_t rpt2 = sizeGm.GetValue(size_n - 2);
            for (int i = 0; i < rpt2; i++)
            {
                for (int j = 0; j < rpt1; j++)
                {
                    yGm.SetValue(i * rpt1 + j, xGm.GetValue(offset + j * step1));
                }
                offset += step2;
            }
        }
        else if (size_n == 3)
        {
            int32_t step2 = strideGm.GetValue(size_n - 2);
            int32_t rpt2 = sizeGm.GetValue(size_n - 2);
            int32_t step3 = strideGm.GetValue(size_n - 3);
            int32_t rpt3 = sizeGm.GetValue(size_n - 3);
            for (int i = 0; i < rpt3; i++)
            {
                int32_t offset2 = offset;
                for (int p = 0; p < rpt2; p++)
                {
                    for (int j = 0; j < rpt1; j++)
                    {
                        yGm.SetValue(i * rpt1 * rpt2 + p * rpt1 + j, xGm.GetValue(offset2 + j * step1));
                    }
                    offset2 += step2;
                }
                offset += step3;
            }
        }

    }

private:
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;
    GlobalTensor<DTYPE_SIZE> sizeGm;
    GlobalTensor<DTYPE_SIZE> strideGm;
};
extern "C" __global__ __aicore__ void as_strided(GM_ADDR x, GM_ADDR size, GM_ADDR stride, GM_ADDR storage_offset, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelAss op;
    op.Init(x, size, stride, storage_offset, y, tiling_data.size_n, tiling_data.blockLength);
} 
